from CvPythonExtensions import *

import GerikSpells
from GerikSpells import GerikSpellBase
from GSGroupPolicies import GSAll as GSGroup
from GSHandlerPolicies import GSTargetting as GSHandler
from GSResultsPolicies import GSMission as GSResults

from GSTargetInfoPolicies import GSSquarePlotTarget
from GSTargetInfoPolicies import GSSpecificUnitTarget

from GSGraphicsPolicies import GSPlot
from GSGraphicsPolicies import GSPieMenu

import GerikSpellsUtils

from CivcraftSCUtils import checkAndRemoveEnergy
from CivcraftSCUtils import showNotEnoughEnergyMsg
from CivcraftSCUtils import getEnergyUsageString

gc = CyGlobalContext()
localtext = CyTranslator()

class ArbiterBase(GerikSpellBase, GSGroup, GSHandler):
    def __init__(self):
        GerikSpellBase.__init__(self)
        GSGroup.__init__(self)
        GSHandler.__init__(self)


        self.m_iArbiterType = gc.getInfoTypeForString("UNIT_ARBITER")
            
    def canUnitCast(self, pCastingUnit, pPlot, bVisible):
        if not pCastingUnit.getUnitType() == self.m_iArbiterType:
            return False

        # Show this for all Arbiters.
        if bVisible:
            return True

        return True


class StasisField(ArbiterBase, GSResults):
    def __init__(self):
        ArbiterBase.__init__(self)

        # Constants
        self.m_iStasisFieldEnergyCost = 100
        self.m_iStasisFieldSize = 3
        self.m_iStasisFieldRange = 5
        self.m_iStasisFieldUnitAttachment = gc.getInfoTypeForString("UNIT_ATTACHMENT_STASIS_FIELD")

        self.m_iResearchNeeded = gc.getInfoTypeForString("PROJECT_PROTOSS_ARBITER_STASIS_FIELD")
        self.m_iResearchBuildingToBuildAt = gc.getProjectInfo(self.m_iResearchNeeded).getBuildingToBuildAt()
        self.m_szNotValidText = localtext.getText( "TXT_KEY_CIVCRAFT_REQUIRES_RESEARCH_AT", (gc.getProjectInfo(self.m_iResearchNeeded).getDescription(),
                                                                                             gc.getBuildingInfo(self.m_iResearchBuildingToBuildAt).getDescription()) )

        # Policies
        pSquarePlotTargetPolicy = GSSquarePlotTarget(self.canTargetPlot, self.m_iStasisFieldSize)
        GSResults.__init__(self, pSquarePlotTargetPolicy)

    def getHelpText(self, pGroup, pPlot):
        szText = getEnergyUsageString(self.m_iStasisFieldEnergyCost)

        bValid = self.canGroupCast(pGroup, pPlot, False)
        if not bValid:
            szText += self.m_szNotValidText
        else:
            szText += gc.getSpellInfo(self.getSpellNumber()).getHelp()

        return szText

    def canUnitCast(self, pCastingUnit, pPlot, bVisible):
        if not ArbiterBase.canUnitCast(self, pCastingUnit, pPlot, bVisible):
            return False

        if bVisible:
            return True

        # Check if they have the research
        if not gc.getPlayer( pCastingUnit.getOwner() ).hasBuiltProject( self.m_iResearchNeeded ):
            return False

        return True

    def canTargetPlot(self, pCastingGroup, pCastingPlot, pTargetPlot, bVisible):
        return True

    def getRangePlotList(self, pCastingPlot):
        return GerikSpellsUtils.getRangePlotDistanceList(pCastingPlot, self.m_iStasisFieldRange)

    def createTargettingStageList(self, bAlt, bShift, bCtrl):
        pPlotTargetPolicy = GSSquarePlotTarget(self.canTargetPlot, self.m_iStasisFieldSize)
        pPlotTargetPolicy.setRangeMakerCallback(self.getRangePlotList)
        pPlotTargetPolicy.setCheckRange(True)

        pPlotGraphicsPolicy = GSPlot()
        pPlotGraphicsPolicy.setShowRangeGraphics(True)
        pPlotGraphicsPolicy.setCursorType("CURSOR_SINGLE_ATTACK")
        return [(pPlotTargetPolicy, pPlotGraphicsPolicy)]

    def isAffected(self, pUnit):
        if pUnit.isInvulnerable():
            return False

        if gc.getUnitInfo(pUnit.getUnitType()).isBuilding():
            return False

        return True

    def startSpellResults(self, pCastingGroup, pTargetInfo):
        bShowNotEnoughEnergy = False
        
        iPlayer = pCastingGroup.getOwner()
        for pUnit in pCastingGroup.getUnits():
            if checkAndRemoveEnergy(pUnit, self.m_iStasisFieldEnergyCost):
                for pPlot in pTargetInfo.getAffectedTargetList(pCastingGroup, pCastingGroup.plot()):
                    for iUnit in range(pPlot.getNumUnits()):
                        pLoopUnit = pPlot.getUnit(iUnit)
                        if self.isAffected(pLoopUnit):
                            pLoopUnit.attachUnitAttachment(iPlayer, self.m_iStasisFieldUnitAttachment)
            else:
                bShowNotEnoughEnergy = True

        if bShowNotEnoughEnergy:
            showNotEnoughEnergyMsg(iPlayer)


class Recall(ArbiterBase, GSResults):
    def __init__(self):
        ArbiterBase.__init__(self)

        # Constants
        self.m_iRecallEnergyCost = 150
        self.m_iRecallSize = 3

        self.m_iResearchNeeded = gc.getInfoTypeForString("PROJECT_PROTOSS_ARBITER_STASIS_FIELD")
        self.m_iResearchBuildingToBuildAt = gc.getProjectInfo(self.m_iResearchNeeded).getBuildingToBuildAt()
        self.m_szNotValidText = localtext.getText( "TXT_KEY_CIVCRAFT_REQUIRES_RESEARCH_AT", (gc.getProjectInfo(self.m_iResearchNeeded).getDescription(),
                                                                                             gc.getBuildingInfo(self.m_iResearchBuildingToBuildAt).getDescription()) )

        # Policies
        pSquarePlotTargetPolicy = GSSquarePlotTarget(self.canTargetPlot, self.m_iRecallSize)
        GSResults.__init__(self, pSquarePlotTargetPolicy)

    def getHelpText(self, pGroup, pPlot):
        szText = getEnergyUsageString(self.m_iRecallEnergyCost)

        bValid = self.canGroupCast(pGroup, pPlot, False)
        if not bValid:
            szText += self.m_szNotValidText
        else:
            szText += gc.getSpellInfo(self.getSpellNumber()).getHelp()

        return szText

    def canUnitCast(self, pCastingUnit, pPlot, bVisible):
        if not ArbiterBase.canUnitCast(self, pCastingUnit, pPlot, bVisible):
            return False

        if bVisible:
            return True

        # Check if they have the research
        if not gc.getPlayer( pCastingUnit.getOwner() ).hasBuiltProject( self.m_iResearchNeeded ):
            return False

        return True

    def canTargetPlot(self, pCastingGroup, pCastingPlot, pTargetPlot, bVisible):
        return True

    def createTargettingStageList(self, bAlt, bShift, bCtrl):
        pPlotTargetPolicy = GSSquarePlotTarget(self.canTargetPlot, self.m_iRecallSize)
        pPlotTargetPolicy.setCheckRange(False)

        pPlotGraphicsPolicy = GSPlot()
        pPlotGraphicsPolicy.setShowRangeGraphics(True)
        pPlotGraphicsPolicy.setCursorType("CURSOR_SINGLE_ATTACK")
        return [(pPlotTargetPolicy, pPlotGraphicsPolicy)]

    def isAffected(self, pUnit):
        if pUnit.isInvulnerable():
            return False

        print pUnit
        print pUnit.getUnitType()
        if gc.getUnitInfo(pUnit.getUnitType()).isBuilding():
            return False

        return True

    def putUnit(self, pCastingPlot, pTargettedPlot, pUnit):
        if not self.isAffected(pUnit):
            return
        
        pPlot = pUnit.plot()
        iDX = pPlot.getX() - pTargettedPlot.getX()
        iDY = pPlot.getY() - pTargettedPlot.getY()

        iResultX = pCastingPlot.getX() + iDX
        iResultY = pCastingPlot.getY() + iDY
        # If we can't move into this plot...
        if not pUnit.canMoveInto(CyMap().plot(iResultX, iResultY), False, False, False):
            # drownUnit(pUnit)    Haha.. maybe in another mod.
            pUnit.setXY(iResultX, iResultY)
            pUnit.jumpToNearestValidPlot()
        else:
            pUnit.setXY(iResultX, iResultY)
        

    def startSpellResults(self, pCastingGroup, pTargetInfo):
        bShowNotEnoughEnergy = False
        pTargettedPlot = pTargetInfo.getTarget()
        
        iPlayer = pCastingGroup.getOwner()
        for pUnit in pCastingGroup.getUnits():
            if checkAndRemoveEnergy(pUnit, self.m_iRecallEnergyCost):
                for pPlot in pTargetInfo.getAffectedTargetList(pCastingGroup, pCastingGroup.plot()):
                    pPlotUnitList = []
                    for iUnit in range(pPlot.getNumUnits()):
                        pPlotUnitList.append(pPlot.getUnit(iUnit))

                    for pLoopUnit in pPlotUnitList:
                        self.putUnit(pCastingGroup.plot(), pTargettedPlot, pLoopUnit)
                            
            else:
                bShowNotEnoughEnergy = True

        if bShowNotEnoughEnergy:
            showNotEnoughEnergyMsg(iPlayer)
